from gensim.models import Word2Vec
from sklearn.decomposition import PCA
import plotly.express as px
from datetime import datetime
import os
import platform
import time


def draw_word2vec_visualization(version:int):
    # 加载Word2Vec模型
    model = Word2Vec.load(f'./models/word2vec-tinymodel-v{version}.model')
    creation_time = os.path.getctime(f'./models/word2vec-tinymodel-v{version}.model')
    creation_time_readable = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(creation_time))
    # 提取需要可视化的词汇和对应的词向量
    words = list(model.wv.index_to_key)[:600]  # 选择前200个常见词汇
    word_vectors = model.wv[words]

    # 使用 PCA 将词向量降维到 2D
    pca = PCA(n_components=2)
    word_vectors_2d = pca.fit_transform(word_vectors)

    # 获取当前时间
    now = datetime.now()

    # 使用 Plotly 绘制图形
    fig = px.scatter(
        x=word_vectors_2d[:, 0], 
        y=word_vectors_2d[:, 1], 
        text=words, 
        title=f'Word2Vec 词向量 2D 可视化'
    )
    fig.update_traces(textposition='top center', marker=dict(size=8, color='LightSkyBlue'))

    # 添加注释，将模型参数信息显示在图表下方
    fig.add_annotation(
        text=(
            f"模型版本 v{version} "
            f"训练时间：{creation_time_readable}<br>"
            f"v_size: {model.vector_size}, "
            f"window: {model.window}, "
            f"min_count: {model.min_count}, "
            f"sg: {model.sg}, "
            f"epochs: {model.epochs}, "
            f"negative: {model.negative}, "
            f"sample: {model.sample}, "
            f"words: {len(model.wv.index_to_key)}"
        ),
        xref="paper",  
        yref="paper",
        x=0.5,  
        y=-0.13,  
        showarrow=False,  
        font=dict(size=12),
    )

    # 显示图表
    fig.show()
